import json
import os
import re

from model_api import qwen3_reranker_api
from tqdm import tqdm

path_dir = os.path.dirname(os.path.abspath(__file__))


def call_rerank_main(qas):
    question, answer = qas
    res = qwen3_reranker_api.call_api(question, answer)
    return res[0]['relevance_score']


if __name__ == '__main__':
    phase5_model = 'glm-4-plus'

    with open(os.path.join(path_dir, "phase5_output", f"{phase5_model}_expand_questions.jsonl"), 'r', encoding='utf-8') as fp:
        lines = fp.readlines()
        ex_questions_data = [json.loads(line) for line in lines]

    raw_qa = [(ex_question_data['question'], ex_question_data['answer']) for ex_question_data in ex_questions_data]
    from concurrent.futures import ProcessPoolExecutor
    with ProcessPoolExecutor(max_workers=10) as executor:
        raw_qa_scores = executor.map(call_rerank_main, raw_qa)
    raw_qa_scores = list(raw_qa_scores)
    print('raw_qa_scores ok')

    flatten_ex_qa = []
    new_question_cnt = 5
    for ex_question_data in ex_questions_data:
        for new_question in ex_question_data['new_questions']:
            flatten_ex_qa.append((new_question, ex_question_data['answer']))

    with ProcessPoolExecutor(max_workers=10) as executor:
        flatten_ex_qa_scores = executor.map(call_rerank_main, flatten_ex_qa)
    flatten_ex_qa_scores = list(flatten_ex_qa_scores)
    ex_qa_scores = [
        flatten_ex_qa_scores[i: i+new_question_cnt]
        for i in range(0, len(flatten_ex_qa_scores), new_question_cnt)
    ]
    assert len(ex_qa_scores) == len(ex_questions_data) == len(raw_qa_scores)

    output_fp = open(os.path.join(path_dir, "phase6_output", "qwen3_reranker_verify_expand_questions.jsonl"), 'w', encoding='utf-8')
    for ex_question_data, raw_qa_score, ex_qa_score in zip(ex_questions_data, raw_qa_scores, ex_qa_scores):
        ex_question_data['score'] = raw_qa_score
        ex_question_data['new_questions_score'] = ex_qa_score
        output_fp.write(json.dumps(ex_question_data, ensure_ascii=False) + '\n')
    output_fp.close()

